# 画图展示目标物体边界框
import os

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.image import imread
from utils import b64decode, generate_timestamp_tuple
import config
import math

# 定义画图函数
INSECT_NAMES = ['Boerner', 'Leconte', 'Linnaeus',
                'acuminatus', 'armandi', 'coleoptera', 'linnaeus']


# 定义画矩形框的函数
def draw_rectangle(currentAxis, bbox, edgecolor='k', facecolor='y', fill=False, linestyle='-'):
    # currentAxis，坐标轴，通过plt.gca()获取
    # bbox，边界框，包含四个数值的list， [x1, y1, x2, y2]
    # edgecolor，边框线条颜色
    # facecolor，填充颜色
    # fill, 是否填充
    # linestype，边框线型
    # patches.Rectangle需要传入左上角坐标、矩形区域的宽度、高度等参数
    rect = patches.Rectangle((bbox[0], bbox[1]), bbox[2] - bbox[0] + 1, bbox[3] - bbox[1] + 1, linewidth=1,
                             edgecolor=edgecolor, facecolor=facecolor, fill=fill, linestyle=linestyle)
    currentAxis.add_patch(rect)


# 保存结果的函数
def save_results(last_data, filename, draw_thresh=0.5):
    last_data['timestamps'].append(generate_timestamp_tuple("save_results", "begin"))
    result = b64decode(last_data['result']['ret'])
    plt.figure(figsize=(10, 10))
    im = imread(filename)
    plt.imshow(im)
    currentAxis = plt.gca()
    colors = ['r', 'g', 'b', 'k', 'y', 'c', 'purple']
    for item in result:
        box = item[2:6]
        label = int(item[0])
        name = INSECT_NAMES[label]
        if item[1] > draw_thresh:
            draw_rectangle(currentAxis, box, edgecolor=colors[label])
            plt.text(box[0], box[1], name, fontsize=12, color=colors[label])
    plt.savefig(os.path.join(config.save_path, os.path.basename(filename)))
    last_data['result'] = {}
    last_data['timestamps'].append(generate_timestamp_tuple("save_results", "done"))
    return last_data
